# 导入必要的模块
from flask import Flask, render_template, request
from flask_cors import CORS
from flask_sqlalchemy import SQLAlchemy
import os
import sys
from datetime import datetime
from werkzeug.utils import secure_filename

# 添加项目根目录到Python路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.append(project_root)

# 导入文档处理相关模块
from document_processor import DocumentProcessor
from document_chunking import SmartDocumentChunker
from embdding_remotealiyun import VectorStore

# 初始化Flask应用
app = Flask(__name__)

# 启用CORS支持
CORS(app)

# 设置模板和静态文件目录
app.template_folder = 'templates'
app.static_folder = 'static'

# 配置数据库
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///database.db'
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False

db = SQLAlchemy(app)

# 设置文件上传目录
UPLOAD_FOLDER = os.path.join(os.path.dirname(__file__), 'uploads')
if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)
print(f'文件上传目录: {UPLOAD_FOLDER}')

# 定义数据库模型
class FileRecord(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    filename = db.Column(db.String(100))
    filepath = db.Column(db.String(200))
    upload_time = db.Column(db.DateTime, default=datetime.now)

# 首页路由 - 显示文件上传页面
@app.route('/')
def index():
    return render_template('index.html')

# 文件上传接口
@app.route('/upload', methods=['POST'])
def upload_file():
    try:
        if 'file' not in request.files:
            return {'error': '未选择文件'}, 400, {'Content-Type': 'application/json'}
        
        file = request.files['file']
        if file.filename == '':
            return {'error': '空文件名'}, 400, {'Content-Type': 'application/json'}
        
        # 安全保存文件
        filename = secure_filename(file.filename)
        save_path = os.path.join(UPLOAD_FOLDER, filename)
        file.save(save_path)
        
        # 记录到数据库
        new_file = FileRecord(filename=filename, filepath=save_path)
        db.session.add(new_file)
        db.session.commit()

        # 异步处理文档
        try:
            # 创建文档处理器实例
            processor = DocumentProcessor()
            chunker = SmartDocumentChunker()
            vector_store = VectorStore()
            print(f'开始处理文档: {filename}')

            # 处理文档
            doc_info = processor.process_document(save_path, filename)
            content = doc_info['content']
            print(f'文档处理完成，内容长度: {len(content)} 字符')

            # 文档分块
            chunks = chunker.chunk_document(content, {'source': filename})
            print(f'文档分块完成，共 {len(chunks)} 块')
            for i, chunk in enumerate(chunks[:3]):
                print(f'块 {i+1} 内容预览: {chunk["content"][:50]}...')

            # 向量化并存储
            vector_store.add_chunks_to_vector_db(chunks, collection_name='rag_docs')
            vector_store.persist()
            print(f'文档向量存储完成，集合名称: rag_docs')
        except Exception as e:
            print(f'文档处理和向量化失败: {str(e)}')
            import traceback
            traceback.print_exc()
            # 不影响文件上传结果，仅记录错误

        return {'status': 'success', 'path': save_path}, 200, {'Content-Type': 'application/json'}
    except Exception as e:
        # 发生异常时返回JSON格式的错误信息
        return {'error': str(e)}, 500, {'Content-Type': 'application/json'}

# 数据库初始化路由
@app.route('/init-db')
def init_db():
    with app.app_context():
        db.create_all()
    return '数据库初始化成功！'

# 启动应用
if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)